#!/usr/bin/env bash
BASE_PATH="$(cd "$(dirname "$0")";pwd)"
function ssh_setup() {
        echo "start config ssh"
        grep -v "[[:space:]]*#" /etc/ssh/sshd_config | grep "PermitRootLogin no"
        if [ $? -eq 0 ]; then
                rm -f /etc/ssh/sshd_config_ancertbak
                cp /etc/ssh/sshd_config /etc/ssh/sshd_config_ancertbak
                sed -i 's/^PermitRootLogin.*/PermitRootLogin yes/g' /etc/ssh/sshd_config
                grep -v "[[:space:]]*#" /etc/ssh/sshd_config | grep "PermitRootLogin yes"
                if [ $? -eq 0 ];then
                      echo "Modify ssh PermitRootLogin yes success"
                else
                      echo "Modify ssh PermitRootLogin yes failure" && return 1
                fi
        fi
        if [ -d /root/.ssh ]; then
                rm -rf /root/.ssh.bak
                mv /root/.ssh /root/.ssh.bak || return 1
                rm -rf /root/.ssh || return 1
        fi
        ansshpkg=$BASE_PATH"/sshpkg"
        scp -r "$ansshpkg" /root/.ssh || return 1
        chmod 700 /root/.ssh
        chmod 600 /root/.ssh/id_rsa
        chmod 644 /root/.ssh/id_rsa.pub /root/.ssh/authorized_keys /root/.ssh/known_hosts
        service sshd start
}

function ssh_restore() {
        echo "start restore sshconfig"
        if [ -f /etc/ssh/sshd_config_ancertbak ]; then
                rm -f /etc/ssh/sshd_config
                mv /etc/ssh/sshd_config_ancertbak /etc/ssh/sshd_config || return 1
                echo "restore sshd_config file success"
        fi
        if [ -d /root/.ssh.bak ]; then
                rm -rf /root/.ssh
                mv /root/.ssh.bak /root/.ssh || return 1
        fi
        service sshd start
}

if [ -z "$1" ]; then
        echo "Usage:"
        echo "    bash utils/sshconf.sh [setup|restore]"
        exit 0
fi

if [ "$1" == "setup" ]; then
        ssh_setup
elif [ "$1" == "restore" ]; then
        ssh_restore
else
        echo "only support setup/restore"
        exit 1
fi
if [ $? -eq 0 ]; then
        echo "ssh config success"
        exit 0
else
        echo "ssh config failed"
        exit 1
fi